﻿using MySoftLab.DapFlashCore.Probes.DapAccess;
using MySoftLab.DapFlashCore.Utils;

namespace MySoftLab.DapFlashCore.CoreSight
{
    public class CortexM
    {
        // CPUID ARCHITECTURE values
        public const byte ARMv6M = 0xC;
        public const byte ARMv7M = 0xF;

        // CPUID PARTNO values
        public const uint ARM_CortexM0 = 0xC20;
        public const uint ARM_CortexM1 = 0xC21;
        public const uint ARM_CortexM3 = 0xC23;
        public const uint ARM_CortexM4 = 0xC24;
        public const uint ARM_CortexM0p = 0xC60;

        // User-friendly names for core types.
        public static Dictionary<uint, string> CORE_TYPE_NAME = new Dictionary<uint, string>()
        {
            { ARM_CortexM0, "Cortex-M0" },
            { ARM_CortexM1, "Cortex-M1" },
            { ARM_CortexM3, "Cortex-M3" },
            { ARM_CortexM4, "Cortex-M4" },
            { ARM_CortexM0p, "Cortex-M0+" }
        };

        // Map from register name to DCRSR register index.
        //
        // The CONTROL, FAULTMASK, BASEPRI, and PRIMASK registers are special in that they share the
        // same DCRSR register index and are returned as a single value. In this dict, these registers
        // have negative values to signal to the register read/write functions that special handling
        // is necessary. The values are the byte number containing the register value, plus 1 and then
        // negated. So -1 means a mask of 0xff, -2 is 0xff00, and so on. The actual DCRSR register index
        // for these combined registers has the key of 'cfbp'.
        public static Dictionary<string, int> CORE_REGISTER = new Dictionary<string, int>()
        {
               {"r0"  , 0},
               {"r1"  , 1},
               {"r2"  , 2},
               {"r3"  , 3},
               {"r4"  , 4},
               {"r5"  , 5},
               {"r6"  , 6},
               {"r7"  , 7},
               {"r8"  , 8},
               {"r9"  , 9},
               {"r10" , 10},
               {"r11" , 11},
               {"r12" , 12},
               {"sp"  , 13},
               {"r13" , 13},
               {"lr"  , 14},
               {"r14" , 14},
               {"pc"  , 15},
               {"r15" , 15},
               {"xpsr", 16},
               {"msp" , 17},
               {"psp" , 18},
               {"cfbp", 20},
               {"control"  ,-4},
               {"faultmask",-3},
               {"basepri"  ,-2},
               {"primask"  ,-1},
               {"fpscr", 33},
               {"s0" , 0x40},
               {"s1" , 0x41},
               {"s2" , 0x42},
               {"s3" , 0x43},
               {"s4" , 0x44},
               {"s5" , 0x45},
               {"s6" , 0x46},
               {"s7" , 0x47},
               {"s8" , 0x48},
               {"s9" , 0x49},
               {"s10", 0x4a},
               {"s11", 0x4b},
               {"s12", 0x4c},
               {"s13", 0x4d},
               {"s14", 0x4e},
               {"s15", 0x4f},
               {"s16", 0x50},
               {"s17", 0x51},
               {"s18", 0x52},
               {"s19", 0x53},
               {"s20", 0x54},
               {"s21", 0x55},
               {"s22", 0x56},
               {"s23", 0x57},
               {"s24", 0x58},
               {"s25", 0x59},
               {"s26", 0x5a},
               {"s27", 0x5b},
               {"s28", 0x5c},
               {"s29", 0x5d},
               {"s30", 0x5e},
               {"s31", 0x5f}
        };

        public const int TARGET_RUNNING = 1;   // Core is executing code.
        public const int TARGET_HALTED = 2;    // Core is halted in debug mode.
        public const int TARGET_RESET = 3;     // Core is being held in reset.
        public const int TARGET_SLEEPING = 4;  // Core is sleeping due to a wfi or wfe instruction.
        public const int TARGET_LOCKUP = 5;    // Core is locked up.

        // CPUID Register
        public const uint CPUID = 0xE000ED00;
        public const uint CPUID_IMPLEMENTER_MASK = 0xff000000;
        public const int CPUID_IMPLEMENTER_POS = 24;
        public const uint CPUID_IMPLEMENTER_ARM = 0x41;
        public const uint CPUID_VARIANT_MASK = 0x00f00000;
        public const int CPUID_VARIANT_POS = 20;
        public const uint CPUID_ARCHITECTURE_MASK = 0x000f0000;
        public const int CPUID_ARCHITECTURE_POS = 16;
        public const uint CPUID_PARTNO_MASK = 0x0000fff0;
        public const int CPUID_PARTNO_POS = 4;
        public const uint CPUID_REVISION_MASK = 0x0000000f;
        public const int CPUID_REVISION_POS = 0;

        public const uint NVIC_AIRCR = 0xE000ED0C;
        public const uint NVIC_AIRCR_VECTKEY = 0x5FA << 16;
        public const uint NVIC_AIRCR_VECTRESET = 1 << 0;
        public const uint NVIC_AIRCR_SYSRESETREQ = 1 << 2;

        // Core Register Selector Register
        public const uint DCRSR = 0xE000EDF4;
        public const uint DCRSR_REGWnR = 1 << 16;
        public const uint DCRSR_REGSEL = 0x1F;

        // Core Register Data Register
        public const uint DCRDR = 0xE000EDF8;

        // Debug Halting Control and Status Register
        public const uint DHCSR = 0xE000EDF0;
        public const uint C_DEBUGEN = 1 << 0;
        public const uint C_HALT = 1 << 1;
        public const uint C_STEP = 1 << 2;
        public const uint C_MASKINTS = 1 << 3;
        public const uint C_SNAPSTALL = 1 << 5;
        public const uint S_REGRDY = 1 << 16;
        public const uint S_HALT = 1 << 17;
        public const uint S_SLEEP = 1 << 18;
        public const uint S_LOCKUP = 1 << 19;
        public const uint S_RETIRE_ST = 1 << 24;
        public const uint S_RESET_ST = 1 << 25;

        // Debug Exception and Monitor Control Register
        public const uint DEMCR = 0xE000EDFC;
        public const uint DEMCR_TRCENA = 1 << 24;
        public const uint DEMCR_VC_HARDERR = 1 << 10;
        public const uint DEMCR_VC_BUSERR = 1 << 8;
        public const uint DEMCR_VC_CORERESET = 1 << 0;

        public const uint DBGKEY = 0xA05F0000;

        // Thumb bit in XPSR.
        public const uint XPSR_THUMB = 0x01000000;

        private readonly DapAccessCmsisDap _link;
        private readonly DebugPort _dp;
        private readonly AccessPort _ap;

        private uint _arch = 0;
        private uint _core = 0;
        private bool _hasFpu = false;

        public AccessPort Ap => _ap;

        public DapAccessCmsisDap Link => _link;

        public CortexM(DapAccessCmsisDap link, DebugPort dp, AccessPort ap)
        {
            _link = link;
            _dp = dp;
            _ap = ap;
        }

        /// <summary>
        /// Read the CPUID register and determine core type.
        /// </summary>
        public void ReadCoreType()
        {
            var cpuid = _ap.Read32(CPUID).Value;

            if ((cpuid & CPUID_IMPLEMENTER_MASK) >> CPUID_IMPLEMENTER_POS != CPUID_IMPLEMENTER_ARM)
            {
                Console.WriteLine("CPU implementer is not ARM!");
            }

            _arch = (cpuid & CPUID_ARCHITECTURE_MASK) >> CPUID_ARCHITECTURE_POS;
            _core = (cpuid & CPUID_PARTNO_MASK) >> CPUID_PARTNO_POS;

            Console.WriteLine($"CPU core is {CORE_TYPE_NAME[_core]}");
        }

        public void Halt()
        {
            _ap.WriteMemory(DHCSR, DBGKEY | C_DEBUGEN | C_HALT);
            _dp.Flush();
        }

        public void Resume()
        {
            if (GetState() != TARGET_HALTED)
            {
                return;
            }

            _ap.WriteMemory(DHCSR, DBGKEY | C_DEBUGEN);
            _dp.Flush();
        }

        public void Reset(bool softwareReset = true)
        {
            // reset a core. After a call to this function, the core is running
            if (softwareReset)
            {
                try
                {
                    _ap.WriteMemory(NVIC_AIRCR, NVIC_AIRCR_VECTKEY | NVIC_AIRCR_SYSRESETREQ);
                    _dp.Flush(); // Without a flush a transfer error can occur
                }
                catch (DAPAccessIntf.TransferError)
                {
                    _dp.Flush();
                }
            }
            else
            {
                _dp.Reset();
            }

            //// # Now wait for the system to come out of reset. Keep reading the DHCSR until
            //// # we get a good response with S_RESET_ST cleared, or we time out.
            //var startTime = DateTime.Now;
            //while (DateTime.Now - startTime < TimeSpan.FromSeconds(2))
            //{
            //    try
            //    {
            //        var dhcsr = _ap.Read32(DHCSR).Value;
            //        if ((dhcsr & S_RESET_ST) == 0)
            //        {
            //            break;
            //        }
            //    }
            //    catch (DAPAccessIntf.TransferError)
            //    {
            //        _dp.Flush();
            //        Thread.Sleep(10);
            //    }
            //}
        }

        private void PostResetCoreAccesibilityTest()
        {
            bool isTimeout = true;
            var startTime = DateTime.Now;

            uint? dhcsr = null;
            while (DateTime.Now - startTime < TimeSpan.FromSeconds(2))
            {
                try
                {
                    dhcsr = _ap.Read32(DHCSR).Value;
                    if ((dhcsr & S_RESET_ST) == 0)
                    {
                        isTimeout = false;
                        break;
                    }
                    Thread.Sleep(10);
                }
                catch (DAPAccessIntf.TransferError)
                {
                    _dp.Flush();
                    Thread.Sleep(10);
                }
            }

            if (isTimeout)
            {
                if (dhcsr == null)
                {
                    Console.WriteLine("Core is not accessible after reset");
                }
                else
                {
                    Console.WriteLine("Core did not come out of reset within timeout");
                }
            }
        }

        public int GetState()
        {
            var dhcsr = _ap.ReadMemory(DHCSR).Value;

            if ((dhcsr & S_RESET_ST) != 0)
            {
                var newDhcsr = _ap.ReadMemory(DHCSR).Value;
                if ((newDhcsr & S_RESET_ST) != 0 && !((newDhcsr & S_RETIRE_ST) != 0))
                {
                    return TARGET_RESET;
                }
            }

            if ((dhcsr & S_LOCKUP) != 0)
            {
                return TARGET_LOCKUP;
            }
            else if ((dhcsr & S_SLEEP) != 0)
            {
                return TARGET_SLEEPING;
            }
            else if ((dhcsr & S_HALT) != 0)
            {
                return TARGET_HALTED;
            }
            else
            {
                return TARGET_RUNNING;
            }
        }

        public bool IsRunning
        {
            get
            {
                return GetState() == TARGET_RUNNING;
            }
        }

        public bool IsHalted
        {
            get
            {
                return GetState() == TARGET_HALTED;
            }
        }

        public static int RegisterNameToIndex(string reg)
        {
            int regId = 0;
            try
            {
                regId = CORE_REGISTER[reg.ToLower()];
            }
            catch (Exception)
            {
                Console.WriteLine($"cannot find {reg} core register");
            }

            return regId;
        }

        public MyDynamicValue ReadCoreRegister(string reg)
        {
            MyDynamicValue result;

            var regIndex = RegisterNameToIndex(reg);

            var regValue = ReadCoreRegisterRaw(reg);

            // Convert int to float.
            if (regIndex >= 0x40)
            {
                result = new MyDynamicValue(CommonUtils.U32BEToFloat32BE(regValue));
            }
            else
            {
                result = new MyDynamicValue(regValue);
            }

            return result;
        }

        public uint ReadCoreRegisterRaw(string reg)
        {
            var vals = ReadCoreRegisterRaw([reg]);
            return vals[0];
        }

        public uint[] ReadCoreRegisterRaw(string[] regList)
        {
            var regIndices = regList.Select(x => RegisterNameToIndex(x)).ToArray();

            // Sanity check register values
            foreach (var reg in regIndices)
            {
                if (!CORE_REGISTER.Values.Contains(reg))
                {
                    throw new Exception($"unknown reg {reg}");
                }
                else if ((reg >= 128 || reg == 33) && !_hasFpu)
                {
                    throw new Exception("attempt to read FPU register without FPU");
                }
            }

            // Begin all reads and writes
            List<DelayResult<uint>> dhcsrCbList = new List<DelayResult<uint>>();
            List<DelayResult<uint>> regCbList = new List<DelayResult<uint>>();

            foreach (var regIdx in regIndices)
            {
                int reg = regIdx;
                if (reg < 0 && reg >= -4)
                {
                    reg = CORE_REGISTER["cfbp"];
                }

                _ap.WriteMemory(DCRSR, (uint)reg);

                // Technically, we need to poll S_REGRDY in DHCSR here before reading DCRDR. But
                // we're running so slow compared to the target that it's not necessary.
                // Read it and assert that S_REGRDY is set
                var dhcsrCb = _ap.ReadMemory(DHCSR, 32, false);
                var regCb = _ap.ReadMemory(DCRDR, 32, false);

                dhcsrCbList.Add(dhcsrCb);
                regCbList.Add(regCb);
            }

            // Read all results
            List<uint> regVals = new List<uint>();

            for (int i = 0; i < regList.Length; i++)
            {
                uint dhcsrVal = 0;
                var action = dhcsrCbList[i].Action;
                if (action != null)
                {
                    dhcsrVal = action();
                }

                if ((dhcsrVal & S_REGRDY) == 0)
                {
                    throw new Exception();
                }

                uint val = 0;
                action = regCbList[i].Action;
                if (action != null)
                {
                    val = action();
                }

                var reg = regIndices[i];
                // Special handling for registers that are combined into a single DCRSR number.
                if (reg < 0 && reg >= -4)
                {
                    val = val >> (-reg - 1) * 8 & 0xff;
                }

                regVals.Add(val);
            }

            return regVals.ToArray();
        }

        public void WriteCoreRegister(string reg, uint data)
        {
            var regIndex = RegisterNameToIndex(reg);

            //// convert float to int
            //if(regIndex >= 0x40)
            //{
            //    data = CommonUtils.Float32beToU32be();
            //}
            WriteCoreRegisterRaw(reg, data);
        }

        public void WriteCoreRegisterRaw(string reg, uint data)
        {
            WriteCoreRegistersRaw([reg], [data]);
        }

        public void WriteCoreRegistersRaw(string[] regList, uint[] dataList)
        {
            // Write one or more core registers
            if (regList.Length != dataList.Length)
            {
                throw new Exception();
            }

            var regIndices = regList.Select(x => RegisterNameToIndex(x)).ToArray();

            // Sanity check register values
            foreach (var reg in regIndices)
            {
                if (!CORE_REGISTER.Values.Contains(reg))
                {
                    throw new Exception($"unknown reg {reg}");
                }
                else if ((reg >= 128 || reg == 33) && !_hasFpu)
                {
                    throw new Exception("attempt to read FPU register without FPU");
                }
            }

            uint specialRegValue = 0;

            // Read special register if it is present in the list
            foreach (var reg in regIndices)
            {
                if (reg < 0 && reg >= -4)
                {
                    specialRegValue = ReadCoreRegister("cfbp").IntValue;
                    break;
                }
            }

            // Write out registers
            List<DelayResult<uint>> dhcsrCbList = new List<DelayResult<uint>>();

            uint data;
            for (int i = 0; i < regList.Length; i++)
            {
                var reg = regIndices[i];
                data = dataList[i];
                if (reg < 0 && reg >= -4)
                {
                    // Mask in the new special register value so we don't modify the other register values that share the same DCRSR number.
                    var shift = (-reg - 1) * 8;
                    var mask = (uint)(0xffffffff ^ 0xff << shift);
                    data = specialRegValue & mask | (data & 0xff) << shift;
                    specialRegValue = data; // update special register for other writes that might be in the list
                    reg = CORE_REGISTER["cfbp"];
                }

                _ap.WriteMemory(DCRDR, data);
                _ap.WriteMemory(DCRSR, (uint)(reg | DCRSR_REGWnR)); // start write transfer

                // Technically, we need to poll S_REGRDY in DHCSR here to ensure the
                // register write has completed.
                // Read it and assert that S_REGRDY is set
                var dhcsrCb = _ap.ReadMemory(DHCSR, 32, false);
                dhcsrCbList.Add(dhcsrCb);
            }

            // Make sure S_REGRDY was set for all register writes
            foreach (var dhcsrCb in dhcsrCbList)
            {
                if (dhcsrCb.Action != null)
                {
                    var dhcsrVal = dhcsrCb.Action();
                    if ((dhcsrVal & S_REGRDY) == 0)
                    {
                        throw new Exception();
                    }
                }
            }
        }

        public void SetTargetState(string state)
        {
            if (state == "PROGRAM")
            {
                ResetStopOnReset(true);
                // Write the thumb bit in case the reset handler points to an ARM address
                WriteCoreRegister("xpsr", 0x1000000);
            }
        }

        public void ResetStopOnReset(bool softwareReset = true)
        {
            // perform a reset and stop the core on the reset handler
            Console.WriteLine("reset stop on reset");

            Halt();

            var demcr = _ap.ReadMemory(DEMCR).Value;
            _ap.WriteMemory(DEMCR, demcr | DEMCR_VC_CORERESET); // enable the vector catch

            Reset(softwareReset);
            while (IsRunning) ;

            _ap.WriteMemory(DEMCR, demcr);
        }

        public void WriteMemoryBlock32(uint addr, uint[] data)
        {
            _ap.WriteMemoryBlock32(addr, data);
        }

        /// <summary>
        /// Read an aligned block of 32-bit words.
        /// </summary>
        /// <param name="addr"></param>
        /// <param name="size"></param>
        /// <returns></returns>
        public uint[] ReadMemoryBlock32(uint addr, int size)
        {
            var data = _ap.ReadMemoryBlock32(addr, size);
            return data;
        }


        private void SetResetCatch()
        {
            // Halt the target.
            Halt();

            // Enable reset vector catch if needed.
            var demcr = _ap.ReadMemory(DEMCR).Value;

            if ((demcr & DEMCR_VC_CORERESET) == 0)
            {
                _ap.WriteMemory(DEMCR, demcr | DEMCR_VC_CORERESET);
            }
        }

        private void ClearResetCatch()
        {
            // Clear VC_CORERESET in DEMCR.
            var demcr = _ap.ReadMemory(DEMCR).Value;
            if ((demcr & DEMCR_VC_CORERESET) != 0)
            {
                _ap.WriteMemory(DEMCR, demcr & ~DEMCR_VC_CORERESET);
            }
        }

        private void PerformReset()
        {
            _ap.WriteMemory(NVIC_AIRCR, NVIC_AIRCR_VECTKEY | NVIC_AIRCR_SYSRESETREQ);
            _dp.Flush(); // Without a flush a transfer error can occur
            Thread.Sleep(100);
        }

        private void InnerReset(bool isHalting = true)
        {
            PerformReset();
            _dp.PostResetRecovery();
            PostResetCoreAccesibilityTest();

            if (!isHalting)
            {
                if (GetState() == TARGET_HALTED)
                {
                    Console.WriteLine("reset: core was halted after non-halting reset; now resuming");
                    Resume();
                }
            }
        }

        /// <summary>
        /// Make sure the thumb bit is set in XPSR in case the reset handler
        /// points to an invalid address. Only do this if the core is actually halted, otherwise we
        /// can't access XPSR.
        /// </summary>
        private void CheckTBit()
        {
            if (GetState() == TARGET_HALTED)
            {
                uint xpsr = ReadCoreRegisterRaw("xpsr");
                if ((xpsr & XPSR_THUMB) == 0)
                {
                    Console.WriteLine("T bit in XPSR is invalid; the vector table may be invalid or corrupt");
                }
            }
        }

        public void ResetAndHalt()
        {
            SetResetCatch();

            InnerReset(true);

            bool isTimeout = true;
            var now = DateTime.Now;
            int state = 0;
            while (DateTime.Now - now < TimeSpan.FromSeconds(2))
            {
                state = GetState();
                if (state != TARGET_RESET && state != TARGET_RUNNING)
                {
                    isTimeout = false;
                    break;
                }
                Thread.Sleep(100);
            }

            if (isTimeout)
            {
                Console.WriteLine($"Timed out waiting for core to halt after reset (state is {state})");
            }

            ClearResetCatch();
            CheckTBit();
        }

        public void UnlockDevice()
        {

        }
    }
}
